package utils

import (
	"fmt"
	"math/rand"
	"strings"
	"time"
)

// VerifyStockCode 校验Stock代码
// 输入: 6位数字的Stock代码
// 返回: Stock交易所代码
// 支持的交易所和代码前缀：
//
// 上交所 (sh): 600, 601, 603, 605, 688(科创板), 900(B股)
//
// 深交所 (sz): 000, 001, 002(中小板), 003, 200(B股), 300(创业板)
//
// 北交所 (bj): 430, 831-839, 8开头(新三板精选层)
func VerifyStockCode(code string) (exchange string, err error) {
	// 去除输入中的空格和特殊字符
	code = strings.TrimSpace(code)

	// 检查长度是否为6
	if len(code) != 6 {
		return "", fmt.Errorf("stock代码长度必须为6位")
	}

	// 检查是否为纯数字
	for _, c := range code {
		if c < '0' || c > '9' {
			return "", fmt.Errorf("stock代码必须为6位数字")
		}
	}

	// 根据代码前两位或三位判断交易所
	switch {
	// 上交所: 600, 601, 603, 605, 688(科创板), 900(B股)
	case strings.HasPrefix(code, "60"),
		strings.HasPrefix(code, "68"),
		strings.HasPrefix(code, "900"):
		return "sh", nil

	// 深交所: 000, 001, 002(中小板), 003, 200(B股), 300(创业板)
	case strings.HasPrefix(code, "00"),
		strings.HasPrefix(code, "200"),
		strings.HasPrefix(code, "30"):
		return "sz", nil

	// 北交所: 430, 831, 832, 833, 834, 835, 836, 837, 838, 839, 8(新三板精选层)
	case strings.HasPrefix(code, "92"),
		strings.HasPrefix(code, "43"),
		strings.HasPrefix(code, "83"):
		return "bj", nil

	default:
		return "", fmt.Errorf("无法识别的Stock代码: %s", code)
	}
}

// CompleteStockCode 补全代码的交易所前缀
func CompleteStockCode(code string) (completeCode string, err error) {
	switch len(code) {
	case 6:
		prefix, err2 := VerifyStockCode(code)
		if err2 != nil {
			return "", err2
		}
		completeCode = prefix + code
	case 8:
		prefix := code[:2]
		if _, ok := map[string]struct{}{"sz": {}, "sh": {}, "bj": {}}[prefix]; !ok {
			err = fmt.Errorf("不合理代码前缀：%v", code)
			return
		}
		completeCode = code
	default:
		err = fmt.Errorf("不合理代码：%v", code)
		return
	}

	return
}

// GetRandomElement 切片中随机取出一个元素
func GetRandomElement[T any](slice []T) T {
	rand.New(rand.NewSource(time.Now().UnixNano()))
	return slice[rand.Intn(len(slice))]
}
